import io
import re
import json
import base64
import asyncio
import time

import pandas as pd
import dataframe_image as dfi
import matplotlib
import matplotlib.pyplot as plt
import uvicorn
import torch
import os
# Windows 环境 Excel COM 导出
import win32com.client as win32

from tqdm import tqdm
from openai import OpenAI
# FastAPI服务构建相关
from fastapi import FastAPI, UploadFile, File, Form
from fastapi.responses import JSONResponse, FileResponse, StreamingResponse
from fastapi.middleware.cors import CORSMiddleware
from typing import List, Optional
# 模型相关
# from transformers import Qwen2_5_VLForConditionalGeneration, AutoProcessor
# from qwen_vl_utils import process_vision_info
# doc文档相关
from docx import Document
from docx.shared import Inches, Pt, Cm
from docx.text.paragraph import Paragraph
from docx.oxml.ns import qn
from docx.enum.text import WD_ALIGN_PARAGRAPH, WD_LINE_SPACING
# 其他相关
from matplotlib.ticker import FormatStrFormatter, PercentFormatter
from tempfile import NamedTemporaryFile

client = OpenAI(
    api_key="sk-8f3aff8d01834f1799127ad52a3d481e", # 配置API-KEY
    base_url="https://dashscope.aliyuncs.com/compatible-mode/v1", # 可修改相关url
)

matplotlib.rcParams['font.sans-serif'] = ['SimHei']
matplotlib.rcParams['axes.unicode_minus'] = False

def strip_number_prefix(txt: str) -> str:
    """
    去掉一、（一）、1. 等各种前缀
    :param txt:
    :return:
    """
    txt = re.sub(r'^[一二三四五六七八九十]+、\s*', '', txt)
    txt = re.sub(r'^[（(]?[一二三四五六七八九十]+[)）]、\s*', '', txt)
    txt = re.sub(r'^\.*\s*\d+\.\s*', '', txt)
    return txt.strip()

def extract_paragraph_groups_by_text(doc_bytes: bytes) -> List[List[Paragraph]]:
    """
    Word 按“此处呈现”分段，保留 Paragraph 对象
    :param doc_bytes:
    :return:
    """
    doc = Document(io.BytesIO(doc_bytes))
    groups: List[List[Paragraph]] = [[]]
    delimiter_re = re.compile(r"^[（(]?此处呈现.*[)）]?$")
    for p in doc.paragraphs:
        txt = p.text.strip()
        if not txt:
            continue
        if delimiter_re.match(txt):
            if groups[-1]:
                groups.append([])
        else:
            groups[-1].append(p)
    return [grp for grp in groups if grp]

def df_to_png_base64(df: pd.DataFrame, pct_cols: Optional[List[str]] = None) -> str:
    """
    DataFrame → 表格 PNG (隐藏索引 & 居中 & 格式化)，并跳过空表
    :param df:
    :param pct_cols:
    :return:
    """
    if df.empty:
        return ""
    num_cols = df.select_dtypes(include="number").columns.tolist()
    if pct_cols is None:
        auto_pct = [c for c in num_cols if df[c].abs().max() <= 1]
    else:
        auto_pct = [c for c in num_cols if c in pct_cols]
    fmt_map = {c: "{:.1%}" if c in auto_pct else "{:.1f}" for c in num_cols}

    styler = (
        df.style
          .format(fmt_map)
          .hide(axis="index")
          .set_table_styles([{"selector": "th", "props": [("text-align", "center")]}])
          .set_properties(**{"text-align": "center"})
    )
    buf = io.BytesIO()
    dfi.export(styler, buf, table_conversion="matplotlib")
    return "data:image/png;base64," + base64.b64encode(buf.getvalue()).decode()

def chart_to_png_base64(
        df: pd.DataFrame,
        title: Optional[str] = None,
        pct_cols: Optional[List[str]] = None
) -> str:
    """
    程序生成图表 PNG (带标题 & 百分比 & 一位小数)
    :param df:
    :param title:
    :param pct_cols:
    :return:
    """
    num_df = df.select_dtypes(include="number")
    if num_df.empty:
        return ""
    if pct_cols is None:
        auto_pct = [c for c in num_df.columns if num_df[c].abs().max() <= 1]
    else:
        auto_pct = [c for c in num_df.columns if c in pct_cols]

    fig, ax = plt.subplots()
    if auto_pct:
        for col in auto_pct:
            line, = ax.plot(num_df.index, num_df[col], marker="o", label=col)
            for x, y in zip(line.get_xdata(), line.get_ydata()):
                ax.text(x, y, f"{y * 100:.1f}%", ha="center", va="bottom", fontsize=8)
        ax.yaxis.set_major_formatter(PercentFormatter(1.0))
    else:
        lines = []
        for col in num_df.columns:
            line, = ax.plot(num_df.index, num_df[col], marker="o", label=col)
            lines.append(line)
        ax.yaxis.set_major_formatter(FormatStrFormatter("%.1f"))
        for line in lines:
            for x, y in zip(line.get_xdata(), line.get_ydata()):
                ax.text(x, y, f"{y:.1f}", ha="center", va="bottom", fontsize=8)

    if title:
        ax.set_title(title, fontproperties=matplotlib.font_manager.FontProperties(family='SimHei'))
    ax.legend(prop={'family': 'SimHei'})
    buf = io.BytesIO()
    fig.savefig(buf, format="png", bbox_inches="tight")
    plt.close(fig)
    return "data:image/png;base64," + base64.b64encode(buf.getvalue()).decode()

def extract_chart_images_from_excel(excel_bytes: bytes) -> dict:
    """
    提取 Excel 中已有图表为 PNG
    :param excel_bytes:
    :return:
    """
    with NamedTemporaryFile(delete=False, suffix=".xlsx") as tmp:
        tmp.write(excel_bytes)
        tmp_path = tmp.name
    chart_map = {}
    excel_app = win32.Dispatch('Excel.Application')
    excel_app.Visible = False
    wb = excel_app.Workbooks.Open(tmp_path)
    try:
        for sht in wb.Sheets:
            name = sht.Name
            try:
                charts = sht.ChartObjects()
                if charts.Count >= 1:
                    img_path = tmp_path.replace('.xlsx', f'_{name}.png')
                    charts.Item(1).Chart.Export(img_path, FilterName='PNG')
                    with open(img_path, 'rb') as f:
                        chart_map[name] = 'data:image/png;base64,' + base64.b64encode(f.read()).decode()
                    os.remove(img_path)
            except Exception:
                pass
    finally:
        wb.Close(False)
        excel_app.Quit()
        os.remove(tmp_path)
    return chart_map

async def collect_payloads(word_bytes: bytes, excel_bytes: bytes, chart_map: dict) -> List[dict]:
    """
    收集所有 payloads 到列表，并提取“title_text”
    :param word_bytes:
    :param excel_bytes:
    :param chart_map:
    :return:
    """
    para_groups = extract_paragraph_groups_by_text(word_bytes)
    sheets = pd.read_excel(io.BytesIO(excel_bytes), sheet_name=None)
    sheet_names = list(sheets.keys())

    # 缓存当前层级标题
    current_title1 = None
    current_title2 = None
    current_title3 = None

    # Word 样式名
    heading1 = ('Heading 1', '标题 1')
    heading2 = ('Heading 2', '标题 2')
    heading4 = ('Heading 4', '标题 4')

    # 纯文本标题前缀正则
    pat1 = re.compile(r"^[一二三四五六七八九十]+、\s*")
    pat2 = re.compile(r"^[（(]?[一二三四五六七八九十]+[)）]、\s*")
    pat3 = re.compile(r"^\.*\s*\d+\.\s*")

    all_payloads: List[dict] = []

    for idx, group in enumerate(para_groups):
        if idx >= len(sheet_names):
            break
        name = sheet_names[idx]
        df = sheets[name]

        # —— 表格与图表图片 ——
        table_b64 = df_to_png_base64(df)

        if name.startswith('图'):
            # 图XXX：优先 Excel 原图，否则程序绘制
            chart_b64 = chart_map.get(name) or chart_to_png_base64(df, title=strip_number_prefix(name))
        elif name.startswith('表'):
            # 表XXX：复制表格图两份
            chart_b64 = table_b64
        else:
            # 其它：先 Excel 图，再程序绘制
            chart_b64 = chart_map.get(name) or chart_to_png_base64(df, title=strip_number_prefix(name))

        # —— 识别段内所有标题及其索引 ——
        found1 = found2 = found4 = None
        i1 = i2 = i4 = None
        for i, p in enumerate(group):
            raw = p.text.strip()
            style = p.style.name
            if style in heading1 or pat1.match(raw):
                found1, i1 = strip_number_prefix(raw), i
            elif style in heading2 or pat2.match(raw):
                found2, i2 = strip_number_prefix(raw), i
            elif style in heading4 or pat3.match(raw):
                found4, i4 = strip_number_prefix(raw), i

        # —— 更新当前标题缓存 ——
        if found1 is not None:
            current_title1 = found1
            current_title2 = found2
            current_title3 = found4
        elif found2 is not None:
            current_title2 = found2
            current_title3 = found4
        elif found4 is not None:
            current_title3 = found4

        # —— 提取 title_text ——
        title_text = 0
        segment = []
        if i2 is not None and i4 is not None and i4 > i2 + 1:
            segment = group[i2+1:i4]
        elif i1 is not None and i2 is not None and i2 > i1 + 1:
            segment = group[i1+1:i2]

        if segment:
            lines = []
            for p in segment:
                t = p.text.strip()
                if not (pat1.match(t) or pat2.match(t) or pat3.match(t)):
                    lines.append(t)
            if lines:
                title_text = "\n".join(lines)

        # —— 构造正文 text —— 跳过所有标题 + title_text 段落 ——
        skip_idxs = set(ix for ix in (i1, i2, i4) if ix is not None)
        for p in segment:
            skip_idxs.add(group.index(p))

        text_lines = []
        for j, p in enumerate(group):
            if j in skip_idxs:
                continue
            t = p.text.strip()
            if pat1.match(t) or pat2.match(t) or pat3.match(t):
                continue
            text_lines.append(t)
        text_block = "\n".join(text_lines)

        # —— 构造 payload ——
        if not (found1 or found2 or found4):
            payload = {
                "title1": 0,
                "title2": 0,
                "title3": 0,
                "title_text": 0
            }
        else:
            payload = {
                "title1": current_title1 or 0,
                "title2": current_title2 or 0,
                "title3": current_title3 or 0,
                "title_text": title_text
            }

        payload.update({
            "text": text_block,
            "sheet": name,
            "table_image": table_b64
        })
        if chart_b64:
            payload["chart_image"] = chart_b64

        all_payloads.append(payload)
        await asyncio.sleep(0)

    return all_payloads

# 使用本地部署模型则打开此部分
# def model_process(input_data: str):
#     """
#     处理传入的照片和模板语句
#     :param input_data:
#     :return:
#     """
#     try:
#         # 保存上传的图片到临时文件
#         # 图表图片
#         temp_file_path1 = f"./image/output1.png"
#
#         # 构建输入消息
#         messages = [
#             {
#                 "role": "user",
#                 "content": [
#                     {"type": "image", "image": temp_file_path1},
#                     {"type": "text", "text": input_data}
#                 ]
#             }
#         ]
#
#         # 处理输入
#         text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
#
#         image_inputs, video_inputs = process_vision_info(messages)
#
#         inputs = processor(
#             text=[text],
#             images=image_inputs,
#             videos=video_inputs,
#             padding=True,
#             return_tensors="pt",
#         )
#         inputs = inputs.to(model.device)
#
#         # 推理生成
#         generated_ids = model.generate(**inputs, max_new_tokens=2048)
#         generated_ids_trimmed = [
#             out_ids[len(in_ids):] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
#         ]
#         # 输出结果
#         output_text = processor.batch_decode(
#             generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
#         )
#
#         # 删除临时文件
#
#         os.remove(temp_file_path1)
#
#         return JSONResponse(content={"result": output_text[0], "code": 200}, status_code=200)
#
#     except Exception as e:
#         return JSONResponse(content={"error": str(e), "code": 500}, status_code=500)


# 模型部分
def change_base64(base64_data1: str, base64_data2: str):
    """
    将图片保存到本地，多个
    :param base64_data1:
    :param base64_data2:
    :return:
    """
    img_data1 = base64.b64decode(base64_data1)
    img_data2 = base64.b64decode(base64_data2)
    # 写入文件
    with open('./image/output1.png', 'wb') as f:
        f.write(img_data1)

    with open('./image/output2.png', 'wb') as f:
        f.write(img_data2)

def change_base64_simple(base64_data1: str):
    """
    将图片保存到本地，单个
    :param base64_data1:
    :return:
    """
    img_data1 = base64.b64decode(base64_data1)
    # 写入文件
    with open('./image/output1.png', 'wb') as f:
        f.write(img_data1)

def encode_image_to_base64(image_path):
    """
    读取本地图片并编码成base64
    :param image_path:
    :return:
    """
    with open(image_path, "rb") as f:
        image_data = f.read()
    return base64.b64encode(image_data).decode('utf-8')

def model_url(input_data: str):
    """
    异步访问模型接口，生成doc文档
    :param input_data:
    :return:
    """
    # 本地图片路径
    local_image_path = "./image/output1.png"  # 把路径换成你的实际本地路径
    image_base64 = encode_image_to_base64(local_image_path)
    time.sleep(5)
    reasoning_content = ""  # 定义完整思考过程
    answer_content = ""  # 定义完整回复
    is_answering = False  # 判断是否结束思考过程并开始回复
    prompt = """角色设定：

你是一位“智能教育报告生成专家”。

任务目标：

请根据用户提供的图表与输出模板生成一份教育分析报告。

执行要求：

尽可能严格遵循用户提供的输出模板格式进行报告内容组织。

在生成报告前，仔细比对图表数据与模板结构，判断是否完全一致。

若图表与模板结构存在不一致，需采取以下处理方式：
a. 识别模板中“没有对应图表数据结果”的部分；
b. 将这些“没有结果内容的部分”从最终报告中删除；
c. 在撰写过程中就应考虑该删除策略，而非等生成完再处理。

保持语言简洁准确，避免生成无关或冗余内容。

输出方式：

最终报告应以条理清晰、结构规范的方式输出，满足模板格式和用户预期。
    """
    try:
        completion = client.chat.completions.create(
            model="qvq-max-latest",
            messages=[
                {
                    "role": "system",
                    "content": [
                        {
                            "type": "text",
                            # "text": "你是一名智能教育报告的专家，请根据图片中的数据结合用户的输出后来替代用户输出模板格式中的XX数据，并严格按照用户输出模板格式输出，并给出分析过程。\n其中中等及以上是水平Ⅱ、Ⅲ、Ⅳ的比例值之和。\n最后输出模板用以下格式：\n输出模板：\n'''\n输出结果\n'''"
                            # "text": "你是一名智能教育报告的专家，请根据用户提供的输出模板进行输出，如果用户的图和模板出现不一致的情况，最后没有结果部分的内容请删除后再输出，并且思考过程尽可能精练简洁。"
                            "text": prompt
                        }
                    ],
                },
                {
                    "role": "user",
                    "content": [
                        {
                            "type": "image_url",
                            "image_url": {
                                "url": f"data:image/jpeg;base64,{image_base64}"
                            },
                        },
                        {"type": "text", "text": input_data},
                    ],
                },
            ],
            stream=True,
            temperature=0.1
        )
        for chunk in completion:
            # 如果chunk.choices为空，则打印usage
            if not chunk.choices:
                print("\nUsage:")
                print(chunk.usage)
            else:
                delta = chunk.choices[0].delta
                # 打印思考过程
                if hasattr(delta, 'reasoning_content') and delta.reasoning_content != None:
                    print(delta.reasoning_content, end='', flush=True)
                    reasoning_content += delta.reasoning_content
                else:
                    # 开始回复
                    if delta.content != "" and is_answering is False:
                        # print("\n" + "=" * 20 + "完整回复" + "=" * 20 + "\n")
                        is_answering = True
                    # 打印回复过程
                    print(delta.content, end='', flush=True)
                    answer_content += delta.content
        return JSONResponse(content={"result": answer_content, "code": 200}, status_code=200)
    except Exception as e:
        print("api处出错")
        os.remove("./image/output1.png")
        return JSONResponse(content={"error": str(e), "code": 500}, status_code=500)

async def model_chat(input_data: str, flag: bool):
    """
    异步访问模型接口，与模型对话
    :param input_data:
    :param flag:
    :return:
    """
    async def event_generator():
        reasoning_content = ""  # 累积思考过程
        answer_content = ""  # 累积回答内容
        is_answering = False  # 是否已开始回答
        flag_reason = 1
        flag_content = 1
        prompt = """角色设定：

你是一位“智能教育报告生成专家”。

任务目标：

请根据用户提供的图表与输出模板生成一份教育分析报告。

执行要求：

尽可能严格遵循用户提供的输出模板格式进行报告内容组织。

在生成报告前，仔细比对图表数据与模板结构，判断是否完全一致。

若图表与模板结构存在不一致，需采取以下处理方式：
a. 识别模板中“没有对应图表数据结果”的部分；
b. 将这些“没有结果内容的部分”从最终报告中删除；
c. 在撰写过程中就应考虑该删除策略，而非等生成完再处理。

保持语言简洁准确，避免生成无关或冗余内容。

输出方式：

最终报告应以条理清晰、结构规范的方式输出，满足模板格式和用户预期。
    """
        if flag:
            local_image_path = "./image/output1.png"
            image_base64 = encode_image_to_base64(local_image_path)
            messages = [
                {
                    "role": "system",
                    "content": [
                        {
                            "type": "text",
                            # "text": "你是一名智能教育报告的专家，请根据用户提供的输出模板进行输出..."
                            "text": prompt
                        }
                    ],
                },
                {
                    "role": "user",
                    "content": [
                        {
                            "type": "image_url",
                            "image_url": {
                                "url": f"data:image/jpeg;base64,{image_base64}"
                            },
                        },
                        {"type": "text", "text": input_data},
                    ],
                },
            ]
        else:
            messages = [
                {
                    "role": "system",
                    "content": [
                        {
                            "type": "text",
                            # "text": "你是一名智能教育报告的专家，请准确回答用户的提问"
                            "text": prompt
                        }
                    ],
                },
                {
                    "role": "user",
                    "content": [
                        {"type": "text", "text": input_data},
                    ],
                },
            ]
        try:
            completion = client.chat.completions.create(
                model="qvq-max-latest",
                messages=messages,
                stream=True,
                temperature=0.1
            )

            for chunk in completion:
                await asyncio.sleep(0)  # 协程让步
                if not chunk.choices:
                    continue
                delta = chunk.choices[0].delta
                if hasattr(delta, 'reasoning_content') and delta.reasoning_content:
                    if flag_reason == 1:
                        flag_reason = 0

                        yield f"[REASON]{delta.reasoning_content}"
                    else:

                        yield f"{delta.reasoning_content}"
                elif delta.content:
                    if not is_answering:
                        is_answering = True
                    if flag_content == 1:
                        flag_content = 0

                        yield f"[ANSWER]{delta.content}"
                    else:

                        yield delta.content

        except Exception as e:
            yield f"[ERROR] {str(e)}"

    return StreamingResponse(event_generator(), media_type="text/plain")

def process_data(input_data):
    """
    对输入的base64图片数据进行处理以及获取到模型返回的结果
    :param table_image:
    :param chart_image:
    :param input_data:
    :return:
    """
    # 放入模型进行处理，得到模型处理后的结果
    # 本地模型部署方式
    # process_data = model_process(input_data).body.decode("utf-8")
    # API方式
    data = model_url(input_data).body.decode("utf-8")
    print(data)
    process_data_analysis = json.loads(data)["result"]
    process_data_code = json.loads(data)["code"]
    if process_data_code == 200:
        return JSONResponse(content={"result": process_data_analysis, "code": 200}, status_code=200)
    else:
        print("process_data处出错")
        return JSONResponse(content={"error": json.loads(data)["error"]}, status_code=500)

def doc_process(doc: Document(), process_data_analysis) -> Document():
    """
    将模型传输过来的数据放入到doc文档中
    :param doc:
    :param process_data_analysis:
    :param title:
    :param level:
    :param flag: 如果为1则是第一次出现需要增加标题
    :return:
    """
    doc = add_paragraph_data(doc, process_data_analysis)

    # 插入图片
    doc.add_picture('./image/output2.png', width=Cm(12.26), height=Cm(7.96))  # 调整插入图片大小
    pic = doc.paragraphs[-1]
    pic.alignment = WD_ALIGN_PARAGRAPH.CENTER
    os.remove("./image/output2.png")
    return doc

def all_process(input_data, doc: Document()) -> Document():
    """
    处理全流程的问题
    :param input_data:
    :param doc:
    :return:
    """
    try:
        data_title3 = process_data(input_data).body.decode("utf-8")
        print(data_title3)
        process_data_analysis = json.loads(data_title3)["result"]
        process_data_code = json.loads(data_title3)["code"]
        if process_data_code != 200:
            raise ValueError(json.loads(data_title3)["error"])
        doc = doc_process(doc, process_data_analysis)
        return doc
    except Exception as e:
        # os.remove("./image/output1.png")
        # os.remove("./image/output2.png")
        print("all_process处出错")
        print("error:", str(e))
        return JSONResponse(content={"error": str(e)}, status_code=500)

def add_paragraph_data(doc: Document(), data) -> Document():
    """
    给doc文档添加段落
    :param doc:
    :param data:
    :return:
    """
    # para = doc.add_paragraph(data)
    # para.paragraph_format.first_line_indent = Cm(0.74)  # 首行缩进0.74cm-两个字符
    # para.alignment = WD_ALIGN_PARAGRAPH.LEFT  # 左对齐
    # para.line_spacing = 1.5  # 行距
    # # 给段落添加一个 run，统一设置字体
    # run = para.add_run()
    # run.text = data
    # run.font.name = '宋体'
    # run._element.rPr.rFonts.set(qn('w:eastAsia'), '宋体')  # 设置中文字体为宋体
    # run.font.size = Pt(12)  # 设置字号为小四号（12磅）
    # return doc
    for para_text in data.split('\n'):
        para_text = para_text.strip()
        if not para_text:
            continue  # 跳过空行
        para = doc.add_paragraph()
        run = para.add_run(para_text)

        # 设置段落格式
        para.paragraph_format.first_line_indent = Cm(0.74)  # 首行缩进0.74cm
        para.alignment = WD_ALIGN_PARAGRAPH.LEFT  # 左对齐
        para.paragraph_format.line_spacing_rule = WD_LINE_SPACING.EXACTLY  # 固定行距
        para.paragraph_format.line_spacing = Pt(28)  # 行距设置为28磅

        # 设置字体样式
        run.font.name = '宋体'
        run._element.rPr.rFonts.set(qn('w:eastAsia'), '宋体')
        run.font.size = Pt(12)
    return doc

def all_process_no_pic(doc: Document(), text) -> Document():
    """
    处理不带图片的情况
    :param doc:
    :param title:
    :param level:
    :param text:
    :return:
    """
    doc = add_paragraph_data(doc, text)

    # 插入图片
    doc.add_picture('./image/output2.png', width=Cm(12.26), height=Cm(7.96))  # 调整插入图片大小
    pic = doc.paragraphs[-1]
    pic.alignment = WD_ALIGN_PARAGRAPH.CENTER
    os.remove("./image/output2.png")
    return doc

def judge(data1, data2):
    """
    判断是否与上次是一样，如果不一样就不用添加,不添加返回的是False
    :param data1:
    :param data2:
    :return:
    """
    if data1 == data2:
        return True
    return False

# FastAPI 应用 & 接口
app = FastAPI()
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# 模型和处理器加载（在服务器启动时加载一次）
# model_dir = "../qwen_vl_test/model"
# model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
#     model_dir,
#     torch_dtype=torch.bfloat16,
#     device_map="auto"
# )
# 需要pip install -U flash-attn --no-build-isolation
# 可启用 flash_attention_2 节省显存和提高速度
# model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
#     model_dir,
#     torch_dtype=torch.bfloat16,
#     attn_implementation="flash_attention_2",
#     device_map="auto",
# )

# processor = AutoProcessor.from_pretrained(model_dir, max_pixels=1280*28*28)

# 可根据像素大小调整
# min_pixels = 256*28*28
# max_pixels = 1280*28*28
# processor = AutoProcessor.from_pretrained(model_dir, min_pixels=min_pixels, max_pixels=max_pixels)

@app.post("/doc/")
async def interface(
    word: UploadFile = File(...),
    excel: UploadFile = File(...)
):
    """
    对输入进行解析，最终返回doc文档
    :param word:
    :param excel:
    :return:
    """
    # 读取参数、解析相关参数
    word_bytes = await word.read()
    excel_bytes = await excel.read()
    chart_map = extract_chart_images_from_excel(excel_bytes)
    payloads = await collect_payloads(word_bytes, excel_bytes, chart_map)
    title1_old = ""
    title2_old = ""
    title3_old = ""
    # 由于最终返回doc文档，要创建一个新的文档
    doc = Document()
    # 标题创建并居中
    title = doc.add_heading('智慧教育示例文档', level=1)
    title.alignment = WD_ALIGN_PARAGRAPH.CENTER
    for payload in tqdm(payloads, desc="文档处理", unit="部分"):
        # 预处理部分
        title1_new = payload["title1"] # 一级标题
        title2_new = payload["title2"] # 二级标题
        title3_new = payload["title3"] # 三级标题
        title_text = payload["title_text"]
        follow_prompt = payload["text"]
        table_image = payload["table_image"]
        chart_image = payload["chart_image"]
        table_image = table_image.split("data:image/png;base64,")[-1]
        chart_image = chart_image.split("data:image/png;base64,")[-1]
        change_base64(table_image, chart_image)

        # 对图片数据进行解析和存储 存在./image/output1.png和./image/output2.png中


        input_data = f"请根据图片中的信息，并严格按照以下格式输出，不需要添加其他话语：\n   {follow_prompt}"
        # 处理读取到的title
        if title1_new == 0:
            doc = all_process(input_data, doc)
        else:
            if not judge(title1_new, title1_old):
                doc.add_heading(title1_new, level=1)
                if title_text != 0:
                    doc = add_paragraph_data(doc, title_text)
                title1_old = title1_new
            if title2_new != 0:
                if not judge(title2_new, title2_old):
                    doc.add_heading(title2_new, level=2)
                    if title_text != 0:
                        doc = add_paragraph_data(doc, title_text)
                    title2_old = title2_new
                # 全为非0的情况
                if title3_new != 0:
                    if not judge(title3_new, title3_old):
                        doc.add_heading(title3_new, level=3)
                        title3_old = title3_new
                    # 标题为三级的情况
                    doc = all_process(input_data, doc)
                # title2存在但是title3为0
                else:
                    # 标题为二级的情况
                    if table_image == "":
                        doc = all_process_no_pic(doc, follow_prompt)
                    else:
                        doc = all_process(input_data, doc)
            else:
                # 标题为一级的情况
                if table_image == "":
                    # 存在没有图片，只用放文字进去的情况
                    doc = all_process_no_pic(doc, follow_prompt)
                else:
                    doc = all_process(input_data, doc)

    # 将文档保存在内存中
    doc_io = io.BytesIO()
    doc.save(doc_io)
    doc_io.seek(0)

    # 返回结果
    # Content-Disposition 是一个 HTTP 头，指示浏览器应如何显示内容。
    # "attachment" 指示浏览器下载文件，而不是直接在浏览器中显示它。
    # filename=example.docx 告诉浏览器下载时文件的默认名称是 example.docx。
    return StreamingResponse(doc_io, media_type='application/vnd.openxmlformats-officedocument.wordprocessingml.document', headers={"Content-Disposition": "attachment; filename=example.docx"})

@app.post("/chat/")
async def chat(
    image: Optional[UploadFile] = File(None),
    question: str = Form()
):
    """
    跟模型进行对话，返回流式输出
    :param image:
    :param question:
    :return:
    """
    if image is None:
        return await model_chat(question, False)
    else:
        # 读取文件二进制内容
        image_bytes = await image.read()

        # 将二进制内容编码为 base64 字符串
        image_base64 = base64.b64encode(image_bytes).decode('utf-8')
        print(image_base64)
        change_base64_simple(image_base64)
        return await model_chat(question, True)



uvicorn.run(app, host="127.0.0.1", port=8000)
